import asyncio
import time
import datetime
import numpy as np
import os
import sys
from bleak import BleakClient
from argparse import ArgumentParser

# MAC 地址
MAC_ADDRESS = '26:EF:60:96:87:3A'

# IMU 特征 UUID（根据我们之前的扫描）
IMU_CONTROL_CHARACTERISTIC_UUID = "0000ae01-0000-1000-8000-00805f9b34fb"
IMU_DATA_CHARACTERISTIC_UUID = "0000ae02-0000-1000-8000-00805f9b34fb"

# 全局变量
limt_count = 0
parse_imu_flage = False
sock_pc = None

# 数据解析比例因子
scaleAccel = 0.00478515625      # 加速度 [-16g~+16g]
scaleQuat = 0.000030517578125   # 四元数 [-1~+1]
scaleAngle = 0.0054931640625    # 角度 [-180~+180]
scaleAngleSpeed = 0.06103515625 # 角速度 [-2000~+2000]
scaleMag = 0.15106201171875     # 磁场 [-4950~+4950]
scaleTemperature = 0.01         # 温度
scaleAirPressure = 0.0002384185791  # 气压 [-2000~+2000]
scaleHeight = 0.0010728836      # 高度 [-9000~+9000]

def parse_imu(buf):
    """解析 IMU 数据"""
    global limt_count
    
    imu_dat = [0.0] * 34
    
    if buf[0] == 0x11:
        ctl = (buf[2] << 8) | buf[1]
        print(f"subscribe tag: 0x{ctl:04x}")
        print(f"ms: {((buf[6]<<24) | (buf[5]<<16) | (buf[4]<<8) | buf[3])}")

        L = 7  # 从第7字节开始解析数据

        # 解析加速度数据
        if (ctl & 0x0001) != 0:
            for i in range(3):
                value = np.short((np.short(buf[L+1]) << 8) | buf[L]) * scaleAccel
                imu_dat[i] = float(value)
                L += 2
                print(f"a{'XYZ'[i]}: {value:.3f}")

        # 保存数据到CSV
        with open("data.csv", "a") as datafile:
            datafile.write(time.strftime("%Y-%m-%d-%H-%M-%S"))
            for i in range(34):
                datafile.write(f",{imu_dat[i]:.6f}")
            datafile.write("\n")

        print(f'limt_count: {limt_count}')
        limt_count += 1

        # 定期重命名文件
        if limt_count % 3530 == 0:
            tim = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
            if os.path.exists("data.csv"):
                os.rename("data.csv", f"f{tim}.csv")
            limt_count = 0
    else:
        print(f"[error] data head not define: 0x{buf[0]:02x}")

def notification_handler(sender, data):
    """处理从 IMU 设备接收到的数据通知"""
    global parse_imu_flage
    if parse_imu_flage:
        print(f"接收到数据，长度: {len(data)} 字节")
        parse_imu(data)
    else:
        print(f"接收到数据但未启用解析，长度: {len(data)} 字节")

async def configure_device(client):
    """配置 IMU 设备"""
    try:
        print("发现的服务和特征值:")
        for service in client.services:
            print(f"服务 UUID: {service.uuid}")
            for char in service.characteristics:
                print(f"  特征值 UUID: {char.uuid}, 句柄: {char.handle}, 属性: {char.properties}")
        
        # 查找控制特征和数据特征
        # 我们需要选择正确的特征值句柄
        control_char = None
        data_char = None
        
        # 遍历所有特征值，选择正确的句柄
        for service in client.services:
            for char in service.characteristics:
                if char.uuid == IMU_CONTROL_CHARACTERISTIC_UUID:
                    # 选择具有 write-without-response 属性的控制特征
                    if 'write-without-response' in char.properties:
                        control_char = char
                        print(f"找到控制特征: UUID={char.uuid}, 句柄={char.handle}")
                
                elif char.uuid == IMU_DATA_CHARACTERISTIC_UUID:
                    # 选择具有 notify 属性的数据特征
                    if 'notify' in char.properties:
                        data_char = char
                        print(f"找到数据特征: UUID={char.uuid}, 句柄={char.handle}")
        
        if not control_char:
            print("未找到控制特征")
            return False
        
        if not data_char:
            print("未找到数据特征")
            return False
        
        # 发送配置命令
        print("开始配置设备...")
        
        # 保持连接命令 - 使用句柄而不是UUID
        print("发送保持连接命令...")
        await client.write_gatt_char(control_char.handle, b')')
        await asyncio.sleep(0.1)

        # 尝试采用蓝牙高速通信特性
        print("发送蓝牙高速通信命令...")
        await client.write_gatt_char(control_char.handle, bytes([0x46]))
        await asyncio.sleep(0.1)

        # 参数设置
        isCompassOn = 0        # 1=使用磁场融合姿态，0=不使用
        barometerFilter = 2
        Cmd_ReportTag = 0x0FFF  # 功能订阅标识

        params = bytearray([0x00] * 11)
        params[0] = 0x12
        params[1] = 5       # 静止状态加速度阀值
        params[2] = 255     # 静止归零速度(单位cm/s) 0:不归零 255:立即归零
        params[3] = 0       # 动态归零速度(单位cm/s) 0:不归零
        params[4] = ((barometerFilter & 3) << 1) | (isCompassOn & 1)
        params[5] = 60      # 数据主动上报的传输帧率[取值0-250HZ], 0表示0.5HZ
        params[6] = 1       # 陀螺仪滤波系数[取值0-2],数值越大越平稳但实时性越差
        params[7] = 3       # 加速计滤波系数[取值0-4],数值越大越平稳但实时性越差
        params[8] = 5       # 磁力计滤波系数[取值0-9],数值越大越平稳但实时性越差
        params[9] = Cmd_ReportTag & 0xff
        params[10] = (Cmd_ReportTag >> 8) & 0xff

        print("发送参数配置...")
        await client.write_gatt_char(control_char.handle, bytes(params))
        await asyncio.sleep(0.1)

        # 主动上报命令
        print("发送主动上报命令...")
        await client.write_gatt_char(control_char.handle, bytes([0x19]))
        await asyncio.sleep(0.1)

        # 启用通知 - 使用句柄
        print("启用数据通知...")
        await client.start_notify(data_char.handle, notification_handler)
        
        print("设备配置完成!")
        return True

    except Exception as e:
        print(f"配置设备时出错: {e}")
        import traceback
        traceback.print_exc()
        return False

async def main():
    """主函数"""
    global parse_imu_flage
    
    # 解析命令行参数
    arg_parser = ArgumentParser(description="IMU600 Connect Demo")
    arg_parser.add_argument('mac_address', nargs='?', default=MAC_ADDRESS, 
                           help="MAC address of device to connect")
    arg_parser.add_argument('--host-ip', help="HOST ip address of device to connect", default=None)
    args = arg_parser.parse_args()
    
    # 设置是否解析数据
    parse_imu_flage = args.host_ip is None
    
    print(f"尝试连接 IMU 设备: {args.mac_address}")
    
    try:
        # 增加连接超时时间
        async with BleakClient(args.mac_address, timeout=30.0) as client:
            print("连接成功!")
            
            # 检查连接状态
            if client.is_connected:
                print("设备已连接，开始配置...")
                
                # 等待服务发现完成（自动进行）
                await asyncio.sleep(2)
                
                # 配置设备
                if await configure_device(client):
                    print("设备配置成功，开始接收数据...")
                    
                    # 启用数据解析
                    parse_imu_flage = True
                    
                    # 保持连接并持续接收数据
                    print("等待数据输入... (按 Ctrl+C 停止)")
                    try:
                        while True:
                            await asyncio.sleep(1)
                    except KeyboardInterrupt:
                        print("用户中断程序")
                else:
                    print("设备配置失败")
            else:
                print("设备未正确连接")

    except Exception as e:
        print(f"连接或通信错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    # 运行主程序
    asyncio.run(main())

